#!/usr/bin/env python3
"""
预检查脚本：验证对比实验的所有组件是否正常
"""

import sys
import os
import yaml
import pickle
from pathlib import Path
sys.path.insert(0, 'd:\\Alpha-SQL')

def check_imports():
    """检查模块导入"""
    try:
        from alphasql.algorithm.evolutionary.evolutionary_solver import EvolutionarySolver
        from alphasql.runner.mcts_runner import MCTSRunner
        from alphasql.config.mcts_config import MCTSConfig
        import yaml
        import pickle
        from pathlib import Path
        print('✓ 所有必要模块导入成功')
        return True
    except Exception as e:
        print(f'✗ 模块导入失败: {e}')
        return False

def check_config():
    """检查配置文件"""
    try:
        config_path = 'config/qwen32b_bird_dev_evolutionary.yaml'
        with open(config_path, 'r') as f:
            config_data = yaml.safe_load(f)
        print('✓ 进化算法配置文件加载成功')
        print(f'  - solver_type: {config_data.get("solver_type")}')
        evolutionary_kwargs = config_data.get('evolutionary_kwargs', {})
        print(f'  - population_size: {evolutionary_kwargs.get("population_size")}')
        print(f'  - generations: {evolutionary_kwargs.get("generations")}')
        print(f'  - mutation_rate: {evolutionary_kwargs.get("mutation_rate")}')
        print(f'  - save_root_dir: {config_data.get("save_root_dir")}')

        # 检查MCTS配置
        mcts_config_path = 'config/qwen32b_bird_dev.yaml'
        with open(mcts_config_path, 'r') as f:
            mcts_config_data = yaml.safe_load(f)
        print('✓ MCTS配置文件加载成功')
        print(f'  - solver_type: {mcts_config_data.get("solver_type", "mcts")}')
        print(f'  - save_root_dir: {mcts_config_data.get("save_root_dir")}')
        return True
    except Exception as e:
        print(f'✗ 配置文件检查失败: {e}')
        return False

def check_tasks():
    """检查任务文件"""
    try:
        tasks_file = Path('data/preprocessed/bird/dev/dev/tasks.pkl')
        if not tasks_file.exists():
            print('✗ 任务文件不存在')
            return False

        with open(tasks_file, 'rb') as f:
            tasks = pickle.load(f)
        print(f'✓ 任务文件加载成功，包含 {len(tasks)} 个任务')

        if len(tasks) > 0:
            task = tasks[0]
            print(f'  - 第一个任务ID: {task.question_id}')
            print(f'  - 数据库ID: {getattr(task, "db_id", "unknown")}')
            print(f'  - 问题预览: {getattr(task, "question", "")[:100]}...')
        return True
    except Exception as e:
        print(f'✗ 任务文件检查失败: {e}')
        return False

def check_directories():
    """检查目录结构"""
    try:
        dirs_to_check = [
            'results',
            'tools'
        ]
        for dir_path in dirs_to_check:
            path = Path(dir_path)
            if path.exists():
                print(f'✓ 目录存在: {dir_path}')
            else:
                print(f'✗ 目录不存在: {dir_path}')
                return False
        return True
    except Exception as e:
        print(f'✗ 目录检查失败: {e}')
        return False

def check_api_key():
    """检查API密钥"""
    api_key = os.environ.get('OPENAI_API_KEY')
    if api_key:
        print(f'✓ API密钥已设置 (长度: {len(api_key)})')
        return True
    else:
        print('✗ OPENAI_API_KEY 环境变量未设置')
        return False

def main():
    print("=== Alpha-SQL 对比实验预检查 ===\n")

    checks = [
        ("模块导入", check_imports),
        ("配置文件", check_config),
        ("任务文件", check_tasks),
        ("目录结构", check_directories),
        ("API密钥", check_api_key),
    ]

    all_passed = True
    for check_name, check_func in checks:
        print(f"检查 {check_name}...")
        if not check_func():
            all_passed = False
        print()

    if all_passed:
        print("🎉 所有检查通过！可以开始对比实验。")
        print("\n建议运行命令:")
        print("  .\\tools\\run_comparison_experiment.ps1 -Tasks 0,1,2")
        print("\n或者直接运行Python脚本:")
        print("  python tools/run_comparison_experiment.py --tasks 0 1 2")
    else:
        print("❌ 部分检查失败，请修复问题后再开始实验。")

    return all_passed

if __name__ == "__main__":
    success = main()
    sys.exit(0 if success else 1)